import yaml
import argparse

import torch

def to_device(m, x):
    """Send tensor into the device of the module.

    Args:
        m : Torch module, or Torch tensor
        x : Torch tensor.

    Returns:
        Tensor: Torch tensor located in the same place as torch module.

    """
    if isinstance(m, torch.nn.Module):
        device = next(m.parameters()).device
    elif isinstance(m, torch.Tensor):
        device = m.device
    else:
        raise TypeError(
            "Expected torch.nn.Module or torch.tensor, " f"bot got: {type(m)}"
        )
    return x.to(device)



def batch_dict_to_device(batch, device):
    result_d = dict()
    for k, v in batch.items():
        if k == 'phrase_id':
            result_d[k] = v
        else:
            result_d[k] = v.to(device)
    return result_d


def load_cfg(cfg_path):
    with open(cfg_path, 'r') as f:
        config = yaml.load(f, Loader=yaml.Loader)
        return config


def update_cfg(cfg, update_vs):
    """
        cfg: dict, configs to update
        update_vs: dict
    """
    for k, v in cfg.items():
        # print(k, v)
        if k in update_vs.keys() and type(update_vs[k]) == type(v):
            cfg[k] = update_vs[k]
        if type(v) == type(cfg) or isinstance(type(v), dict):
            update_cfg(v, update_vs)


def save_cfg(cfg, cfg_path):
    with open(cfg_path, 'w') as f:
        yaml.dump(cfg, f, Dumper=yaml.Dumper)


class ArgumentParserWithUnkown(argparse.ArgumentParser):

    def parse_to_dict(self, args=None, namespace=None):
        result = dict()
        args, argv = self.parse_known_args(args, namespace)
        # flags = vars(flags)
        if args is not None:
            result.update(vars(args))
        result.update(self.parse_argv(argv))
        return args, result

    @staticmethod
    def parse_argv(argv):
        argv_dict = dict()
        if argv is None:
            return argv_dict
        for ele in argv:
            ele = ele.strip('-')
            eles = ele.split('=')
            k = eles[0]
            v = eles[-1]
            if v.lower() == 'true':
                v = True
            elif v.lower() == 'false':
                v = False
            else:
                try:
                    v = int(v)
                except ValueError:
                    try:
                        v = float(v)
                    except ValueError:
                        pass
            argv_dict[k] = v
        return argv_dict
